// "Copyright [2021] <Copyright Shengkai Lin>"

// The central controller is responsible to manage the new
// tasks from the task generator, alloc gpus, routes, etc for
// the tasks, and manage the critical resources.

#ifndef NETWORK_CONTROLLER_NETWORK_CONTROLLER_H_
#define NETWORK_CONTROLLER_NETWORK_CONTROLLER_H_

#include <vector>
#include <map>
#include "ai_sim.h"
#include "network_controller/topo_manager.h"
#include "network_controller/routing/route.h"

namespace ai_sim {

class TaskGenerator;

class NetworkController : public omnetpp::cSimpleModule{
 private:
    ClusterType clustertype;
    TopoManager *topos;
    TaskGenerator *tasgen;

 public:
    // The following redefined virtual function holds the algorithm.
    void initialize() override;
    void handleMessage(omnetpp::cMessage *msg) override;
    void finish() override;


    void gen_route(taskboard&, RouteAlgorithm*);
    // Used when we need to alloc a batch of routes
    // to multi switches
    void allocRoutes(const map<int, RouteVec>&);

    // Alloc tasks from alloc gpu tp alloc flow and routes
    void alloc_taskroutes(taskboard&);

    void flow_ctrl_msg_send(int, const vector<Flow>&, Task&);

};
}  // namespace ai_sim

#endif  // NETWORK_CONTROLLER_NETWORK_CONTROLLER_H_
